import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim.python.slim.nets import resnet_v1
# from tensorflow.summary import FileWriter
from tensorflow.python.framework import graph_util

Pretrained_model_dir = "/home/dolly/checkpoints/resnet_v1_101_2016_08_28/resnet_v1_101.ckpt"
outputPath = "/home/dolly/checkpoints/resnet_v1_101_2016_08_28/resnet_v1_101.pb"

height, width = 224, 224
X = tf.placeholder(tf.float32, [None, height, width, 3])
# 读取网络
with slim.arg_scope(resnet_v1.resnet_arg_scope()):
    logits, end_points = resnet_v1.resnet_v1_101(X, num_classes=1000, is_training=False)


with tf.Session() as sess:
    # 先初始化所有变量，避免有些变量未读取而产生错误
    init = tf.global_variables_initializer()
    sess.run(init)
    #加载预训练模型
    print('Loading model check point from {:s}'.format(Pretrained_model_dir))

    #这里的exclusions是不需要读取预训练模型中的Logits,因为默认的类别数目是1000，当你的类别数目不是1000的时候，如果还要读取的话，就会报错
    # exclusions = ['InceptionV3/Logits',
    #               'InceptionV3/AuxLogits']
    exclusions = []
    #创建一个列表，包含除了exclusions之外所有需要读取的变量
    inception_except_logits = slim.get_variables_to_restore(exclude=exclusions)

    #建立一个从预训练模型checkpoint中读取上述列表中的相应变量的参数的函数
    init_fn = slim.assign_from_checkpoint_fn(Pretrained_model_dir, inception_except_logits, ignore_missing_vars=False)
    #运行该函数
    init_fn(sess)
    print('Loaded.')

    output_node_names = "save/restore_all"

    builder = tf.saved_model.builder.SavedModelBuilder(outputPath)
    inputs = {'input0': tf.saved_model.utils.build_tensor_info(X)}
    outputs = {'output0': tf.saved_model.utils.build_tensor_info(tf.get_default_graph().get_tensor_by_name(output_node_names))}
    method_name = tf.saved_model.signature_constants.PREDICT_METHOD_NAME

    my_signature = tf.saved_model.signature_def_utils.build_signature_def(inputs, outputs, method_name)
    builder.add_meta_graph_and_variables(sess, ['MODEL_TRAINING'], signature_def_map={'my_signature': my_signature})
    builder.add_meta_graph(['MODEL_SERVING'], signature_def_map={'my_signature': my_signature})
    builder.save()

    # output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化，将变量值固定
    #     sess=sess,
    #     input_graph_def=sess.graph_def,  # 等于:sess.graph_def
    #     output_node_names=output_node_names.split(","))  # 如果有多个输出节点，以逗号隔开
    #
    # with tf.io.gfile.GFile(outputPath, "wb") as f:  # 保存模型
    #     f.write(output_graph_def.SerializeToString())  # 序列化输出
    #     print("%d ops in the final graph." % len(output_graph_def.node))  # 得到当前图有几个操作节点